#include "BvhProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "Input/Input.h"

void BvhProcess::OnInit()
{
	mRayConsts.position[0] = glm::vec4(0.f, 0.f, 0.f, 1.0);
	mRayConsts.position[1] = glm::vec4(1.f, 1.f, 1.f, 1.0);
	mRayConsts.color = glm::vec4(1.f, 1.f, 0.f, 1.f);

	mLocalEntitys = mGameScene->GetEntitysByType<LocalEntity, Model>();

	mSelectColor = glm::vec4(0.f, 1.f, 1.f, 1.f);
	mUnSelectColor = glm::vec4(1.f, 0.f, 0.f, 1.f);

	mLastMouseX = 0;
	mLastMouseY = 0;

	mUseSAH = true;

	mBvhDepth = 10;

	mCaptureRay = true;

	mShowHit = false;

	for (uint32 i = 0; i < mLocalEntitys.size(); ++i)
	{
		Model& model = mLocalEntitys[i]->GetComponent<Model>();
		glm::mat4 modelMat = mLocalEntitys[i]->GetTransform().GetWorldMatrix();

		for (uint32 j = 0; j < model.mIndices.size(); j += 3)
		{
			const uint32 i_1 = model.mIndices[j];
			const uint32 i_2 = model.mIndices[j + 1];
			const uint32 i_3 = model.mIndices[j + 2];

			const glm::vec3 p_1 = modelMat * glm::vec4(model.mVertexs[i_1].Position, 1.f);
			const glm::vec3 p_2 = modelMat * glm::vec4(model.mVertexs[i_2].Position, 1.f);
			const glm::vec3 p_3 = modelMat * glm::vec4(model.mVertexs[i_3].Position, 1.f);

			mTriangles.push_back(Triangle(p_1, p_2, p_3));
		}
	}

	mBvhTree = BuildBVH(mTriangles, 0, mTriangles.size() - 1, 1000);

	mSahBvhTree = BuildBVHWithSAH(mTriangles, 0, mTriangles.size() - 1, 1000);

	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;

			mUboBuffer = CustomBuffer::Builder()
				.SetData(&uboData)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			uint32 indexData[24] = { 0, 1, 1, 2, 2, 3, 3, 0, 4, 5, 5, 6, 6, 7, 7, 4, 0, 4, 1, 5, 2, 6, 3, 7 };
			mIndexBuffer = IndexBuffer::Builder()
				.SetData(&indexData)
				.SetSize(sizeof(uint32) * 24)
				.SetIndexCount(24)
				.Build();

			mBvhBuffer = BuildBVHBuffer(mBvhTree);

			mSahBvhBuffer = BuildBVHBuffer(mSahBvhTree);

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<BvhVertexShader>()
				.PushShader<BvhFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mDescriptor->SetValue("ubo", mUboBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();

			RefCountPtr<RHITexture> colorRenderTarget = GetTexture("LocalProcess", "ColorRenderTarget");

			RefCountPtr<RHITexture> depthRenderTarget = GetTexture("LocalProcess", "DepthRenderTarget");

			RefCountPtr<RHIRenderPass> renderPass = GetRenderPass("LocalProcess", "RenderPass");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(colorRenderTarget)
				.SetDepthRenderTarget(depthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetLineWidth(1.f)
				.SetRenderPass(renderPass->GetRenderPass())
				.Build();




			mRayDescriptor = RHIDescriptor::Builder()
				.PushShader<RayVertexShader>()
				.PushShader<RayFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mRayDescriptor->SetValue("ubo", mUboBuffer);
			mRayDescriptorMask = mRayDescriptor->AllocateDescriptor();

			mRayPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mRayDescriptor)
				.AddColorRenderTarget(colorRenderTarget)
				.SetDepthRenderTarget(depthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetLineWidth(1.f)
				.SetRenderPass(renderPass->GetRenderPass())
				.Build();
		});
}

void BvhProcess::OnTick(float deltaTime)
{
	if (Input::Get().IsKeyDown(Keyboard::P))
	{
		mCaptureRay = true;
	}
	if (Input::Get().IsKeyDown(Keyboard::O))
	{
		mCaptureRay = false;
	}

	if (mCaptureRay)
	{
		if (Input::Get().LeftButton())
		{
			if (Input::Get().MouseX() != mLastMouseX || Input::Get().MouseY() != mLastMouseY)
			{
				mLastMouseX = Input::Get().MouseX();
				mLastMouseY = Input::Get().MouseY();
				Ray ray = mGameScene->GetCamera().GetScreenRay((float)mLastMouseX / GRenderState->GetWidth(), (float)mLastMouseY / GRenderState->GetHeight());

				mRayConsts.position[0] = glm::vec4(ray.mOrigin, 1.0);
				mRayConsts.position[1] = glm::vec4((ray.mOrigin + ray.mDirection * 100.f), 1.0);
				mRayConsts.color = glm::vec4(1.f, 1.f, 0.f, 1.f);

				if (mUseSAH)
				{
					mHitResult = FindHitNode(ray, mSahBvhTree);
				}
				else
				{
					mHitResult = FindHitNode(ray, mBvhTree);
				}

				if (mHitResult.triangle != nullptr && mHitResult.node != nullptr)
				{
					LOGI(L"screen %d, %d hit yes", Input::Get().MouseX(), Input::Get().MouseY());
				}
				else
				{
					LOGI(L"screen %d, %d hit no", Input::Get().MouseX(), Input::Get().MouseY());
				}
			}
		}
	}




	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;
			mUboBuffer->WriteMemory(&uboData);

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

			VkBuffer indexBuffer = mIndexBuffer->GetBuffer();
			vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

			if (mUseSAH)
			{
				DrawBVH(commandBuffer, mSahBvhBuffer, mBvhDepth);
			}
			else
			{
				DrawBVH(commandBuffer, mBvhBuffer, mBvhDepth);
			}

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mRayPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mRayPipeline->GetPipelineLayout(), 0, 1, &mRayDescriptorMask.GetDescriptorSet(), 0, nullptr);

			vkCmdPushConstants(commandBuffer, mRayPipeline->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(RayConsts), &mRayConsts);

			vkCmdDraw(commandBuffer, 2, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void BvhProcess::OnResize()
{

}

void BvhProcess::OnExit()
{
	FreeBvh(mBvhTree);

	AsyncTask(GRenderThreadId, [this]() {
		FreeBVHBuffer(mBvhBuffer);
		FreeBVHBuffer(mSahBvhBuffer);
		});
}

BvhProcess::BVHBuffer* BvhProcess::BuildBVHBuffer(BVHNode* node)
{
	if (node == nullptr)
		return nullptr;

	BVHBuffer* buffer = NewObject<BVHBuffer>();
	buffer->node = node;

	const float minX = node->aa.x;
	const float minY = node->aa.y;
	const float minZ = node->aa.z;

	const float maxX = node->bb.x;
	const float maxY = node->bb.y;
	const float maxZ = node->bb.z;

	std::array<glm::vec3, 8> vertexs = {
		glm::vec3(minX, minY, minZ),
		glm::vec3(minX, maxY, minZ),
		glm::vec3(maxX, maxY, minZ),
		glm::vec3(maxX, minY, minZ),
		glm::vec3(minX, minY, maxZ),
		glm::vec3(minX, maxY, maxZ),
		glm::vec3(maxX, maxY, maxZ),
		glm::vec3(maxX, minY, maxZ)
	};

	buffer->buffer = VertexBuffer::Builder()
		.SetData(vertexs.data())
		.SetSize(sizeof(glm::vec3) * 8)
		.Build();

	buffer->left = BuildBVHBuffer(node->left);
	buffer->right = BuildBVHBuffer(node->right);

	return buffer;
}

void BvhProcess::DrawBVH(VkCommandBuffer commandBuffer, BVHBuffer* buffer, int32 depth)
{
	if (buffer == nullptr)
		return;

	if (depth <= 0)
		return;

	if (mShowHit && buffer->node != mHitResult.node)
	{
		DrawBVH(commandBuffer, buffer->left, depth - 1);
		DrawBVH(commandBuffer, buffer->right, depth - 1);

		return;
	}

	if (buffer->node == mHitResult.node)
	{
		vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(glm::vec4), &mSelectColor);
	}
	else
	{
		vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(glm::vec4), &mUnSelectColor);
	}

	VkBuffer vertexBuffer = buffer->buffer->GetBuffer();
	VkDeviceSize offsets = 0;
	vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

	vkCmdDrawIndexed(commandBuffer, mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);

	DrawBVH(commandBuffer, buffer->left, depth - 1);
	DrawBVH(commandBuffer, buffer->right, depth - 1);

}

void BvhProcess::FreeBVHBuffer(BVHBuffer* buffer)
{
	if (buffer == nullptr)
		return;

	FreeBVHBuffer(buffer->left);
	FreeBVHBuffer(buffer->right);

	DeleteObject<BVHBuffer>(buffer);
}

HitResult BvhProcess::FindHitNode(Ray& ray, BVHNode* node)
{
	if (node == nullptr)
		return HitResult();

	if (node->n > 0)
	{
		return HitNodeResult(mTriangles, ray, node);
	}

	bool isHitLeft = IsHitNode(ray, node->left);
	bool isHitRight = IsHitNode(ray, node->right);

	HitResult r1, r2;
	if (isHitLeft) r1 = FindHitNode(ray, node->left);
	if (isHitRight) r2 = FindHitNode(ray, node->right);

	return r1.distance < r2.distance ? r1 : r2;
}
